from pymongo import MongoClient,errors
import json
from bson.objectid import ObjectId


# mongodb 连接配置
mongodb_url = 'mongodb://localhost:27017/'
db_name = 'test_database'
collect_name = 'restaurants'


# python 模拟mongoimport
## TODO 待完善 模拟失败
def import_json_to_mongodb(json_file_path, database_name, collection_name):
    # 连接到MongoDB
    client = MongoClient('mongodb://localhost:27017/')
    db = client[database_name]
    collection = db[collection_name]

    # 检查集合是否为空
    if collection.find_one() is None:
        # 读取JSON文件
        with open(json_file_path, 'r') as file:
            for line in file:
                # 解析每一行的JSON数据
                document = json.loads(line.strip())

                # 如果文档中包含'_id'，检查它是否有效
                if '_id' in document:
                    try:
                        # 尝试将'_id'转换为ObjectId
                        document['_id'] = ObjectId(document['_id'])
                    except Exception as e:
                        # 如果转换失败，可以选择记录错误或处理它
                        print(f"Invalid _id '{document['_id']}' in document: {document}")
                        # 选择是否删除'_id'字段
                        # del document['_id']

                # 插入到MongoDB集合中
                try:
                    collection.insert_one(document)
                except errors.WriteError as we:
                    # 如果插入时仍然出现问题，记录错误
                    print(f"Failed to insert document: {document}. Error: {we}")

        # 打印确认信息
        print(f"Imported {collection.count_documents({})} documents into the {collection_name} collection.")
    else:
        print(f"The {collection_name} collection already contains data, skipping import.")

    # 关闭连接
    client.close()

if __name__ == "__main__":
    import_json_to_mongodb('restaurants.json',
                           'test_database',
                           'restaurants')
